#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "types.h"
#include "state_vector.h"

//------------------------------------------------------------------------------

// State Vector memory allocation

sv_t *sv_alloc() {

  sv_t *sv = calloc(1, sizeof(sv_t));

  sv->type = calloc(params->Nr*params->Nt*params->Nz,     sizeof(int));
  sv->vr   = calloc((params->Nr+1)*params->Nt*params->Nz, sizeof(double));
  sv->vt   = calloc(params->Nr*params->Nt*params->Nz,     sizeof(double));
  sv->vz   = calloc(params->Nr*params->Nt*(params->Nz+1), sizeof(double));
  sv->p    = calloc(params->Nr*params->Nt*params->Nz,     sizeof(double));

  return sv;

}

//------------------------------------------------------------------------------

void sv_free(sv_t *sv) {

  free(sv->type);
  free(sv->vr);
  free(sv->vt);
  free(sv->vz);
  free(sv->p);
  free(sv);

}

//------------------------------------------------------------------------------

// Creates a copy of sv

sv_t *sv_dup(sv_t *sv) {

  sv_t *sv_new = calloc(1, sizeof(sv_t));

  sv_new->type = memcpy(calloc(params->Nr*params->Nt*params->Nz,     sizeof(int)),    sv->type, params->Nr*params->Nt*params->Nz     * sizeof(int));
  sv_new->vr   = memcpy(calloc((params->Nr+1)*params->Nt*params->Nz, sizeof(double)), sv->vr,   (params->Nr+1)*params->Nt*params->Nz * sizeof(double)); 
  sv_new->vt   = memcpy(calloc(params->Nr*params->Nt*params->Nz,     sizeof(double)), sv->vt,   params->Nr*params->Nt*params->Nz     * sizeof(double));
  sv_new->vz   = memcpy(calloc(params->Nr*params->Nt*(params->Nz+1), sizeof(double)), sv->vz,   params->Nr*params->Nt*(params->Nz+1) * sizeof(double));
  sv_new->p    = memcpy(calloc(params->Nr*params->Nt*params->Nz,     sizeof(double)), sv->p ,   params->Nr*params->Nt*params->Nz     * sizeof(double));

  return sv_new;

}

//------------------------------------------------------------------------------

// z = alpha * X + beta * Y

void sv_axpby(sv_t *z, double alpha, sv_t *x, double beta, sv_t *y) {

  int i;

  for ( i = 0 ; i < (params->Nr+1)*params->Nt*params->Nz ; i++ ) {

	z->vr[i] = alpha * x->vr[i] + beta * y->vr[i];

  }

  for ( i = 0 ; i < params->Nr*params->Nt*params->Nz ; i++ ) {

	z->vt[i] = alpha * x->vt[i] + beta * y->vt[i];

	//	z->p[i] = alpha * x->p[i] + beta * y->p[i];

  }

  for ( i = 0 ; i < params->Nr*params->Nt*(params->Nz+1) ; i++ ) {

	z->vz[i] = alpha * x->vz[i] + beta * y->vz[i];

  }

}

//------------------------------------------------------------------------------

void sv_fwrite(sv_t *y, FILE *fp) {

  fwrite(y->type, sizeof(int),    params->Nr*params->Nt*params->Nz, fp); 
  fwrite(y->vr,   sizeof(double), (params->Nr+1)*params->Nt*params->Nz, fp); 
  fwrite(y->vt,   sizeof(double), params->Nr*params->Nt*params->Nz, fp); 
  fwrite(y->vz,   sizeof(double), params->Nr*params->Nt*(params->Nz+1), fp); 
  fwrite(y->p,    sizeof(double), params->Nr*params->Nt*params->Nz, fp); 

}

//------------------------------------------------------------------------------

void sv_fread(sv_t *y, FILE *fp) {

  if (
	  (int) fread(y->type, sizeof(int),    params->Nr*params->Nt*params->Nz, fp)     != params->Nr*params->Nt*params->Nz     || 
	  (int) fread(y->vr,   sizeof(double), (params->Nr+1)*params->Nt*params->Nz, fp) != (params->Nr+1)*params->Nt*params->Nz ||  
	  (int) fread(y->vt,   sizeof(double), params->Nr*params->Nt*params->Nz, fp)     != params->Nr*params->Nt*params->Nz     || 
	  (int) fread(y->vz,   sizeof(double), params->Nr*params->Nt*(params->Nz+1), fp) != params->Nr*params->Nt*(params->Nz+1) ||
	  (int) fread(y->p,    sizeof(double), params->Nr*params->Nt*params->Nz, fp)     != params->Nr*params->Nt*params->Nz     
	 ) {

	printf("sv_fread error\n");
	exit(0);

  }

}

//------------------------------------------------------------------------------

int sv_get_type(sv_t *y, int i, int j, int k) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( i < 0 || i >= params->Nr || k < 0 || k >= params->Nz ) {

	return CELL_BOUNDARY;

  }

  s = i + params->Nr * ( j % params->Nt + params->Nt * k );

  return y->type[s];

}

//------------------------------------------------------------------------------

double sv_get_vr(sv_t *y, int i, int j, int k) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i-1, j, k) == CELL_FULL ) {

	s = i + (params->Nr+1) * ( j % params->Nt + params->Nt * k );

	return y->vr[s];

  }

  if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i-1, j, k) == CELL_BOUNDARY ) {

	return 0.0;

  }

  if ( sv_get_type(y, i, j, k) == CELL_BOUNDARY && sv_get_type(y, i-1, j, k) == CELL_FULL ) {

	return 0.0;

  }

  if ( sv_get_type(y, i, j, k) == CELL_BOUNDARY && sv_get_type(y, i-1, j, k) == CELL_BOUNDARY ) {

	if ( sv_get_type(y, i, j, k+1) == CELL_FULL && sv_get_type(y, i-1, j, k+1) == CELL_FULL ) {

	  s = i + (params->Nr+1) * ( j % params->Nt + params->Nt * (k+1) );

	  return -(y->vr[s]);

	}

	if ( sv_get_type(y, i, j, k-1) == CELL_FULL && sv_get_type(y, i-1, j, k-1) == CELL_FULL ) {

	  s = i + (params->Nr+1) * ( j % params->Nt + params->Nt * (k-1) );

	  return -(y->vr[s]);

	}

  }

  printf("sv_get_vr error %i %i %i\n", i, j, k);
  exit(0);

}

//------------------------------------------------------------------------------

double sv_get_vt(sv_t *y, int i, int j, int k) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i, j-1, k) == CELL_FULL ) {

	s = i + params->Nr * ( j % params->Nt + params->Nt * k );

	return y->vt[s];

  }

  if ( sv_get_type(y, i, j, k) == CELL_BOUNDARY && sv_get_type(y, i, j-1, k) == CELL_BOUNDARY ) {

	if ( sv_get_type(y, i, j, k+1) == CELL_FULL && sv_get_type(y, i, j-1, k+1) == CELL_FULL ) {

	  s = i + params->Nr * ( j % params->Nt + params->Nt * (k+1) );

	  return -(y->vt[s]);

	}

	if ( sv_get_type(y, i, j, k-1) == CELL_FULL && sv_get_type(y, i, j-1, k-1) == CELL_FULL ) {

	  s = i + params->Nr * ( j % params->Nt + params->Nt * (k-1) );

	  return -(y->vt[s]);

	}

	if ( sv_get_type(y, i+1, j, k) == CELL_FULL && sv_get_type(y, i+1, j-1, k) == CELL_FULL ) {

	  s = (i+1) + params->Nr * ( j % params->Nt + params->Nt * k );

	  return -(y->vt[s]);

	}

	if ( sv_get_type(y, i-1, j, k) == CELL_FULL && sv_get_type(y, i-1, j-1, k) == CELL_FULL ) {

	  s = (i-1) + params->Nr * ( j % params->Nt + params->Nt * k );

	  return -(y->vt[s]);

	}

  }

  printf("sv_get_vt error\n");
  exit(0);

}

//------------------------------------------------------------------------------

double sv_get_vz(sv_t *y, int i, int j, int k) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i, j, k-1) == CELL_FULL ) {

	s = i + params->Nr * ( j % params->Nt + params->Nt * k );

	return y->vz[s];

  }

  if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i, j, k-1) == CELL_BOUNDARY ) {

	return 0.0;

  }

  if ( sv_get_type(y, i, j, k) == CELL_BOUNDARY && sv_get_type(y, i, j, k-1) == CELL_FULL ) {

	return 0.0;

  }

  if ( sv_get_type(y, i, j, k) == CELL_BOUNDARY && sv_get_type(y, i, j, k-1) == CELL_BOUNDARY ) {

	if ( sv_get_type(y, i+1, j, k) == CELL_FULL && sv_get_type(y, i+1, j, k-1) == CELL_FULL ) {

	  s = (i+1) + params->Nr * ( j % params->Nt + params->Nt * k );

	  return -(y->vz[s]);

	}

	if ( sv_get_type(y, i-1, j, k) == CELL_FULL && sv_get_type(y, i-1, j, k-1) == CELL_FULL ) {

	  s = (i-1) + params->Nr * ( j % params->Nt + params->Nt * k );

	  return -(y->vz[s]);

	}

	// Not used yet

	//	if ( sv_get_type(y, i, j+1, k) == CELL_FULL && sv_get_type(y, i, j+1, k-1) == CELL_FULL ) {
	//
	//	  s = i + params->Nr * ( (j+1) % params->Nt + params->Nt * k );
	//
	//	  return -(y->vz[s]);
	//
	//	}
	//
	//	if ( sv_get_type(y, i, j-1, k) == CELL_FULL && sv_get_type(y, i, j-1, k-1) == CELL_FULL ) {
	//
	//	  s = i + params->Nr * ( (j-1) % params->Nt + params->Nt * k );
	//
	//	  return -(y->vz[s]);
	//
	//	}

  }

  printf("sv_get_vz error\n");
  exit(0);

}

//------------------------------------------------------------------------------

void sv_set_vr(sv_t *z, int i, int j, int k, double value) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(z, i, j, k) == CELL_FULL && sv_get_type(z, i-1, j, k) == CELL_FULL ) {

	s = i + (params->Nr+1) * ( j % params->Nt + params->Nt * k );

	z->vr[s] = value;

  }

}

//------------------------------------------------------------------------------

void sv_set_vt(sv_t *z, int i, int j, int k, double value) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(z, i, j, k) == CELL_FULL && sv_get_type(z, i, j-1, k) == CELL_FULL ) {

	s = i + params->Nr * ( j % params->Nt + params->Nt * k );

	z->vt[s] = value;

  }

}

//------------------------------------------------------------------------------

void sv_set_vz(sv_t *z, int i, int j, int k, double value) {

  int s;

  if ( j < 0 ) j += params->Nt;

  if ( sv_get_type(z, i, j, k) == CELL_FULL && sv_get_type(z, i, j, k-1) == CELL_FULL ) {

	s = i + params->Nr * ( j % params->Nt + params->Nt * k );

	z->vz[s] = value;

  }

}

//------------------------------------------------------------------------------

